import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 配置设备（优先用 GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ------------------- 1. 数据加载与预处理 ------------------- #
transform = transforms.Compose([
    transforms.Resize((32, 32)),  # LeNet-5 要求输入 32×32
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 单通道归一化
])

# 加载 FashionMNIST 数据集
train_dataset = torchvision.datasets.FashionMNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = torchvision.datasets.FashionMNIST(
    root='./data', train=False, download=True, transform=transform
)

# 数据加载器（batch 训练）
batch_size = 64
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False
)


# ------------------- 2. 构建 LeNet-5 模型 ------------------- #
class LeNet5(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet5, self).__init__()
        # 卷积层：32×32→28×28→14×14→7×7
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=0)  # 1 通道→6 通道
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, padding=0)  # 6 通道→16 通道
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 全连接层：7×7×16→120→84→10
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 注意：32→28→14→7，这里原 LeNet 是 5×5？
        # 修正：FashionMNIST 是 32×32 输入，经两次池化后是  (32-5+1)=28 → 池化 14 → (14-5+1)=10 → 池化 5 → 所以是 5×5
        # 重新计算：
        # 输入 32×32 → conv1(5x5, 无 padding) → 28×28 → pool → 14×14
        # conv2(5x5, 无 padding) → 10×10 → pool → 5×5
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        # 卷积层 + 池化层
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        x = torch.relu(self.conv2(x))
        x = self.pool2(x)

        # 展平
        x = x.view(-1, 16 * 5 * 5)

        # 全连接层
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# 初始化模型
model = LeNet5().to(device)

# ------------------- 3. 训练与评估 ------------------- #
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10  # 可调整
train_losses = []
train_accuracies = []
test_losses = []
test_accuracies = []

for epoch in range(num_epochs):
    # 训练阶段
    model.train()
    train_correct = 0
    train_total = 0
    train_running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_total += labels.size(0)
        train_correct += (predicted == labels).sum().item()

    train_epoch_loss = train_running_loss / len(train_loader)
    train_epoch_accuracy = 100 * train_correct / train_total
    train_losses.append(train_epoch_loss)
    train_accuracies.append(train_epoch_accuracy)

    # 测试阶段
    model.eval()
    test_correct = 0
    test_total = 0
    test_running_loss = 0.0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            test_running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            test_total += labels.size(0)
            test_correct += (predicted == labels).sum().item()

    test_epoch_loss = test_running_loss / len(test_loader)
    test_epoch_accuracy = 100 * test_correct / test_total
    test_losses.append(test_epoch_loss)
    test_accuracies.append(test_epoch_accuracy)

    # 打印日志
    print(f'Epoch {epoch + 1}/{num_epochs}')
    print(f'Train Loss: {train_epoch_loss:.4f} | Train Acc: {train_epoch_accuracy:.2f}%')
    print(f'Test Loss: {test_epoch_loss:.4f} | Test Acc: {test_epoch_accuracy:.2f}%')
    print('-' * 30)

# ------------------- 4. 可视化曲线 ------------------- #
plt.figure(figsize=(12, 5))

# Loss 曲线
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training vs Test Loss')
plt.legend()

# Accuracy 曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.title('Training vs Test Accuracy')
plt.legend()

plt.show()

# ------------------- 5. 可视化预测结果 ------------------- #
model.eval()
with torch.no_grad():
    # 取一个 batch 的测试数据
    images, labels = next(iter(test_loader))
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)

    # 转换为 numpy（方便可视化）
    images = images.cpu().numpy()
    labels = labels.cpu().numpy()
    predicted = predicted.cpu().numpy()

    # 还原归一化的图像（方便看）
    images = images[:, 0, :, :] * 0.5 + 0.5  # 单通道还原

    # 展示 10 个样本
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    axes = axes.ravel()

    for i in range(10):
        axes[i].imshow(images[i], cmap='gray')
        axes[i].set_title(f'True: {labels[i]}\nPred: {predicted[i]}')
        axes[i].axis('off')

    plt.show()